import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gamma
from sklearn.cluster import DBSCAN
import matplotlib.patches as patches
import pickle


def matern_cluster_process(
    center, radius, num_clusters, lambda_p, mean_cluster_size, gamma_shape
):
    """
    模拟Matern簇过程

    :param center: 圆形区域的中心坐标 (x0, y0)
    :param radius: 圆形区域的半径
    :param num_clusters: 母点（簇中心）的数量
    :param lambda_p: 母点过程的强度
    :param mean_cluster_size: 平均簇大小
    :param gamma_shape: 用于生成簇半径的伽马分布的形状参数
    :return: 模拟得到的所有点的坐标数组，簇中心点的x坐标数组，簇中心点的y坐标数组
    """
    # 生成母点（簇中心）的坐标
    parent_points_x = np.random.uniform(
        center[0] - radius, center[0] + radius, num_clusters
    )
    parent_points_y = np.random.uniform(
        center[1] - radius, center[1] + radius, num_clusters
    )

    all_points = []

    for i in range(num_clusters):
        # 生成簇半径
        cluster_radius = gamma.rvs(gamma_shape)

        # 生成簇内点的数量，假设服从泊松分布
        num_points_in_cluster = np.random.poisson(mean_cluster_size)

        # 生成簇内点的坐标
        points_in_cluster_x = np.random.uniform(
            parent_points_x[i] - cluster_radius,
            parent_points_x[i] + cluster_radius,
            num_points_in_cluster,
        )
        points_in_cluster_y = np.random.uniform(
            parent_points_y[i] - cluster_radius,
            parent_points_y[i] + cluster_radius,
            num_points_in_cluster,
        )

        # 筛选出在圆形区域内的点
        valid_indices = np.where(
            (points_in_cluster_x - center[0]) ** 2
            + (points_in_cluster_y - center[1]) ** 2
            <= radius**2
        )
        valid_points_x = points_in_cluster_x[valid_indices]
        valid_points_y = points_in_cluster_y[valid_indices]

        all_points.extend([(x, y) for x, y in zip(valid_points_x, valid_points_y)])

    return np.array(all_points), parent_points_x, parent_points_y


# 设置参数
center = (0, 0)  # 圆形区域中心坐标
radius = 20  # 圆形区域半径
num_clusters = 20  # 母点（簇中心）数量
lambda_p = 0.5  # 母点过程的强度
mean_cluster_size = 8  # 平均簇大小
gamma_shape = 2  # 用于生成簇半径的伽马分布的形状参数

# 模拟Matern簇过程
points, parent_points_x, parent_points_y = matern_cluster_process(
    center, radius, num_clusters, lambda_p, mean_cluster_size, gamma_shape
)
# 将生成的点数据保存到文件中（使用pickle模块）
with open("matern_cluster_data.pkl", "wb") as f:
    pickle.dump((points, parent_points_x, parent_points_y), f)

# 创建一个包含两个子图的窗口
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 8))

# 获取离坐标原点最近的两个点的距离，以此距离作为半径画圆（修改后的关键逻辑）
distances = np.sqrt((points[:, 0] - center[0]) ** 2 + (points[:, 1] - center[1]) ** 2)
sorted_indices = np.argsort(distances)
radius_w = distances[
    sorted_indices[1]
]  # 取第二个最近的距离作为半径（索引从0开始，所以取1）
circle1 = plt.Circle(center, radius_w, fill=False, color="black")
ax1.add_patch(circle1)
circle2 = plt.Circle(center, radius_w, fill=False, color="black")
ax2.add_patch(circle2)

# 找出圆内的点
in_circle = (points[:, 0] - center[0]) ** 2 + (
    points[:, 1] - center[1]
) ** 2 <= radius_w**2

# 在第一个子图中绘制模拟Matern簇过程的图形
ax1.scatter(points[:, 0], points[:, 1], label="All Points")
parent_points = np.array(
    [(parent_points_x[i], parent_points_y[i]) for i in range(num_clusters)]
)
ax1.scatter(
    parent_points[:, 0],
    parent_points[:, 1],
    marker="^",
    s=50,
    c="green",
    label="Cluster Centers",
)
ax1.set_xlim(center[0] - radius, center[0] + radius)
ax1.set_ylim(center[1] - radius, center[1] + radius)
ax1.set_aspect("equal", adjustable="box")
ax1.set_title("Matern Cluster Process")
ax1.set_xlabel("X-axis")
ax1.set_ylabel("Y-axis")
ax1.legend()

# 剔除圆内的点，获取用于聚类的点
points_for_clustering = points[~in_circle]

# 使用DBSCAN进行聚类
dbscan = DBSCAN(eps=2, min_samples=2)  # 这里的参数可根据实际情况调整
clusters = dbscan.fit_predict(points_for_clustering)

# 提取不同簇的点和噪声点，注意要根据points_for_clustering的索引对应回原points的索引来处理
unique_clusters = np.unique(clusters)

cluster_sets_for_draw = {cluster: [] for cluster in unique_clusters}

cluster_sets = {
    cluster: {"square_indices": [], "no_square_indices": []}
    for cluster in unique_clusters
}
# 存储用于绘图的点
for i, cluster_label in enumerate(clusters):
    if cluster_label != -1:
        # 根据在points_for_clustering中的索引，找到在原points中的对应点
        index_in_original = np.where(~in_circle)[0][i]
        cluster_sets_for_draw[cluster_label].append(points[index_in_original])
    else:
        # 噪声点同样要找到在原points中的对应索引
        noise_indices = np.where(clusters == -1)[0]
        noise_indices_in_original = np.where(~in_circle)[0][noise_indices]
        for idx in noise_indices_in_original:
            cluster_sets_for_draw[-1].append(points[idx])

# 用于存储 points_for_clustering 中每个点在 points 中的原始索引
original_indices = np.where(~in_circle)[0]
for i, cluster_label in enumerate(clusters):
    if cluster_label != -1:
        index_in_original = original_indices[i]
        # 以0.2的概率决定是否为带正方形的点
        if np.random.rand() < 0.2:
            cluster_sets[cluster_label]["square_indices"].append(index_in_original)
        else:
            cluster_sets[cluster_label]["no_square_indices"].append(index_in_original)
    else:
        cluster_sets[-1]["no_square_indices"].append(original_indices[i])


for cluster_label, cluster_points in cluster_sets_for_draw.items():
    if cluster_label != -1:
        ax2.scatter(
            [p[0] for p in cluster_points],
            [p[1] for p in cluster_points],
            label=f"Cluster {cluster_label}",
        )
    else:
        ax2.scatter(
            [p[0] for p in cluster_points],
            [p[1] for p in cluster_points],
            label="Noise Points",
            marker="x",
            c="gray",
        )
parent_points = np.array(
    [(parent_points_x[i], parent_points_y[i]) for i in range(num_clusters)]
)
"""
ax2.scatter(
    parent_points[:, 0],
    parent_points[:, 1],
    marker="^",
    s=50,
    c="green",
    label="Cluster Centers",
)
"""
ax2.set_xlim(center[0] - radius, center[0] + radius)
ax2.set_ylim(center[1] - radius, center[1] + radius)
ax2.set_aspect("equal", adjustable="box")
ax2.set_title("Matern Cluster Process with DBSCAN Clustering")
ax2.set_xlabel("X-axis")
ax2.set_ylabel("Y-axis")
# ax2.legend(loc="upper right", fontsize="small")

# 标注willie
ax1.scatter(0, 0, marker="+", s=100, c="r", label="Willie")
ax2.scatter(0, 0, marker="+", s=100, c="r", label="Willie")


# 将圆内的点变红
ax1.scatter(points[in_circle, 0], points[in_circle, 1], c="red")
ax2.scatter(points[in_circle, 0], points[in_circle, 1], c="red")


# 为每个簇绘制圆（新添加的逻辑）
for cluster_label, cluster_points in cluster_sets.items():
    if cluster_label != -1:
        # 处理带有正方形的点
        square_indices = cluster_points["square_indices"]
        if len(square_indices) > 0:
            square_points = points[square_indices]
            center_x = np.mean(square_points[:, 0])
            center_y = np.mean(square_points[:, 1])
            center_tuple = (center_x, center_y)
            distances = np.sqrt(
                (square_points[:, 0] - center_x) ** 2
                + (square_points[:, 1] - center_y) ** 2
            )
            radius = np.max(distances)
            circle1 = plt.Circle(center_tuple, radius, fill=False, color="blue")
            ax1.add_patch(circle1)
            circle2 = plt.Circle(center_tuple, radius, fill=False, color="blue")
            ax2.add_patch(circle2)
        # 处理不带正方形的点
        no_square_indices = cluster_points["no_square_indices"]
        if len(no_square_indices) > 0:
            no_square_points = points[no_square_indices]

            # no_square_points = points[~square_indices]
            center_x = np.mean(no_square_points[:, 0])
            center_y = np.mean(no_square_points[:, 1])
            center_tuple = (center_x, center_y)
            distances = np.sqrt(
                (no_square_points[:, 0] - center_x) ** 2
                + (no_square_points[:, 1] - center_y) ** 2
            )
            radius = np.max(distances)
            circle1 = plt.Circle(center_tuple, radius, fill=False, color="orange")
            ax1.add_patch(circle1)
            circle2 = plt.Circle(center_tuple, radius, fill=False, color="orange")
            ax2.add_patch(circle2)


# 为 square_indices 对应的点画正方形
for cluster_label, cluster_points in cluster_sets.items():
    if cluster_label != -1:
        square_indices = cluster_points["square_indices"]
        if len(square_indices) > 0:
            square_points = points[square_indices]
            for x, y in square_points:
                rect = patches.Rectangle(
                    (x - 0.5, y - 0.5),
                    1,
                    1,
                    linewidth=1,
                    edgecolor="blue",
                    facecolor="none",
                )
                ax1.add_patch(rect)
                rect = patches.Rectangle(
                    (x - 0.5, y - 0.5),
                    1,
                    1,
                    linewidth=1,
                    edgecolor="blue",
                    facecolor="none",
                )
                ax2.add_patch(rect)

plt.show()
